package pedro.encoder;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.os.Build;
import android.os.Handler;
import android.os.HandlerThread;
import android.util.Log;

import androidx.annotation.NonNull;
import androidx.annotation.RequiresApi;


import java.nio.ByteBuffer;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;

import pedro.encoder.utils.CodecUtil;

/**
 * Created by pedro on 18/09/19.
 */
public abstract class BaseEncoder implements EncoderCallback {
    public enum Status {
        PREPARED, RUNNING, RELEASED
    }

    private static final String TAG = "BaseEncoder";
    private MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
    private HandlerThread handlerThread;
    private MediaCodec.Callback callback;
    private long oldTimeStamp = 0L;

    protected Status status;
    protected BlockingQueue<Frame> queue = new ArrayBlockingQueue<>(80);
    protected MediaCodec codec;
    protected long presentTimeUs;
    protected volatile boolean running = false;
    protected boolean isBufferMode = true;
    protected CodecUtil.Force force = CodecUtil.Force.FIRST_COMPATIBLE_FOUND;

    public void restart() {
        start(false);
        initCodec();
    }

    public void start() {
        start(true);
        initCodec();
    }

    private void initCodec() {
        handlerThread = new HandlerThread(TAG);
        handlerThread.start();
        Handler handler = new Handler(handlerThread.getLooper());
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M) {
            createAsyncCallback();
            codec.setCallback(callback, handler);
            codec.start();
        } else {
            codec.start();
            handler.post(new Runnable() {
                @Override
                public void run() {
                    while (running) {
                        try {
                            getDataFromEncoder();
                        } catch (IllegalStateException e) {
                            Log.i(TAG, "Encoding error", e);
                        }
                    }
                }
            });
        }
        running = true;
        status = Status.RUNNING;
    }

    public Status getStatus() {
        return status;
    }

    public abstract void start(boolean resetTs);

    protected abstract void stopImp();

    protected abstract void reset();

    protected void fixTimeStamp(MediaCodec.BufferInfo info) {
        if (oldTimeStamp > info.presentationTimeUs) {
            info.presentationTimeUs = oldTimeStamp;
        } else {
            oldTimeStamp = info.presentationTimeUs;
        }
    }

    public void stop() {
        running = false;
        stopImp();
        if (handlerThread != null) {
            if (handlerThread.getLooper() != null) {
                if (handlerThread.getLooper().getThread() != null) {
                    handlerThread.getLooper().getThread().interrupt();
                }
                handlerThread.getLooper().quit();
            }
            handlerThread.quit();
            if (codec != null) {
                try {
                    codec.flush();
                } catch (IllegalStateException ignored) {
                }
            }
            //wait for thread to die for 500ms.
            try {
                handlerThread.getLooper().getThread().join(500);
            } catch (Exception ignored) {
            }
        }
        queue.clear();
        queue = new ArrayBlockingQueue<>(80);
        try {
            codec.stop();
            codec.release();
            codec = null;
        } catch (IllegalStateException | NullPointerException e) {
            codec = null;
        }
        oldTimeStamp = 0L;
        status = Status.RELEASED;
    }

    protected abstract MediaCodecInfo chooseEncoder(String mime);

    protected void getDataFromEncoder() throws IllegalStateException {
        if (isBufferMode) {
            int inBufferIndex = codec.dequeueInputBuffer(0);
            if (inBufferIndex >= 0) {
                inputAvailable(codec, inBufferIndex);
            }
        }
        while (running) {
            int outBufferIndex = codec.dequeueOutputBuffer(bufferInfo, 0);
            if (outBufferIndex == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
                MediaFormat mediaFormat = codec.getOutputFormat();
                formatChanged(codec, mediaFormat);
            } else if (outBufferIndex >= 0) {
                outputAvailable(codec, outBufferIndex, bufferInfo);
            } else {
                break;
            }
        }
    }

    protected abstract Frame getInputFrame() throws InterruptedException;

    private void processInput(@NonNull ByteBuffer byteBuffer, @NonNull MediaCodec mediaCodec,
                              int inBufferIndex) throws IllegalStateException {
        try {
            Frame frame = getInputFrame();
            while (frame == null) frame = getInputFrame();
            byteBuffer.clear();
            int size = Math.max(frame.getSize(), 0);
            byteBuffer.put(frame.getBuffer(), frame.getOffset(), size);
            long pts = System.nanoTime() / 1000 - presentTimeUs;
            mediaCodec.queueInputBuffer(inBufferIndex, 0, size, pts, 0);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        } catch (NullPointerException e) {
            Log.i(TAG, "Encoding error", e);
        }
    }

    protected abstract void checkBuffer(@NonNull ByteBuffer byteBuffer,
                                        @NonNull MediaCodec.BufferInfo bufferInfo);

    protected abstract void sendBuffer(@NonNull ByteBuffer byteBuffer,
                                       @NonNull MediaCodec.BufferInfo bufferInfo);

    private void processOutput(@NonNull ByteBuffer byteBuffer, @NonNull MediaCodec mediaCodec,
                               int outBufferIndex, @NonNull MediaCodec.BufferInfo bufferInfo) throws IllegalStateException {
        checkBuffer(byteBuffer, bufferInfo);
        sendBuffer(byteBuffer, bufferInfo);
        mediaCodec.releaseOutputBuffer(outBufferIndex, false);
    }

    public void setForce(CodecUtil.Force force) {
        this.force = force;
    }

    public boolean isBufferMode() {
        return isBufferMode;
    }

    public boolean isRunning() {
        return running;
    }

    @Override
    public void inputAvailable(@NonNull MediaCodec mediaCodec, int inBufferIndex)
            throws IllegalStateException {
        ByteBuffer byteBuffer;
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.LOLLIPOP) {
            byteBuffer = mediaCodec.getInputBuffer(inBufferIndex);
        } else {
            byteBuffer = mediaCodec.getInputBuffers()[inBufferIndex];
        }
        processInput(byteBuffer, mediaCodec, inBufferIndex);
    }

    @Override
    public void outputAvailable(@NonNull MediaCodec mediaCodec, int outBufferIndex,
                                @NonNull MediaCodec.BufferInfo bufferInfo) throws IllegalStateException {
        ByteBuffer byteBuffer;
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.LOLLIPOP) {
            byteBuffer = mediaCodec.getOutputBuffer(outBufferIndex);
        } else {
            byteBuffer = mediaCodec.getOutputBuffers()[outBufferIndex];
        }
        processOutput(byteBuffer, mediaCodec, outBufferIndex, bufferInfo);
    }

    @RequiresApi(api = Build.VERSION_CODES.M)
    private void createAsyncCallback() {
        callback = new MediaCodec.Callback() {
            @Override
            public void onInputBufferAvailable(@NonNull MediaCodec mediaCodec, int inBufferIndex) {
                try {
                    inputAvailable(mediaCodec, inBufferIndex);
                } catch (IllegalStateException e) {
                    Log.i(TAG, "Encoding error", e);
                }
            }

            @Override
            public void onOutputBufferAvailable(@NonNull MediaCodec mediaCodec, int outBufferIndex,
                                                @NonNull MediaCodec.BufferInfo bufferInfo) {
                try {
                    outputAvailable(mediaCodec, outBufferIndex, bufferInfo);
                } catch (IllegalStateException e) {
                    Log.i(TAG, "Encoding error", e);
                }
            }

            @Override
            public void onError(@NonNull MediaCodec mediaCodec, @NonNull MediaCodec.CodecException e) {
                Log.e(TAG, "Error", e);
            }

            @Override
            public void onOutputFormatChanged(@NonNull MediaCodec mediaCodec,
                                              @NonNull MediaFormat mediaFormat) {
                formatChanged(mediaCodec, mediaFormat);
            }
        };
    }
}
